from flask import Flask, render_template, request
import pandas as pd
import joblib

app = Flask(__name__)

# 加载模型
random_forest_model = joblib.load('random_forest_model.pkl')
gradient_boosting_model = joblib.load('gradient_boosting_model.pkl')
svm_model = joblib.load('svm_model.pkl')
mlp_model = joblib.load('mlp_model.pkl')

# 获取商品主键列表
df = pd.read_csv('杭州市.csv', encoding='gbk')
vegetable_types = df['商品主键'].unique().tolist()


@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        # 获取用户输入
        vegetable = request.form.get('vegetable')
        month = int(request.form.get('month'))
        day = int(request.form.get('day'))

        # 创建输入数据
        input_data = pd.DataFrame(columns=random_forest_model.feature_names_in_)
        input_data.loc[0] = 0
        input_data['月份'] = month
        input_data['日'] = day
        input_col = f'商品主键_{vegetable}'
        if input_col in input_data.columns:
            input_data[input_col] = 1

        # 进行预测
        rf_prediction = random_forest_model.predict(input_data)[0]
        gb_prediction = gradient_boosting_model.predict(input_data)[0]
        svm_prediction = svm_model.predict(input_data)[0]
        mlp_prediction = mlp_model.predict(input_data)[0]
        return render_template('index.html', vegetable=vegetable, month=month, day=day,
                               rf_prediction=rf_prediction, gb_prediction=gb_prediction,
                               svm_prediction=svm_prediction, mlp_prediction=mlp_prediction,
                               vegetable_types=vegetable_types)
    return render_template('index.html', vegetable_types=vegetable_types)


if __name__ == '__main__':
    app.run(debug=True)
